#这是最小二乘拟合的通用程序

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import warnings
import sys

warnings.filterwarnings("ignore")#忽略警告以显示高次拟合
#最小二乘拟合
class LSF():
    def __init__(self, dispersed=True, polynomial=True, x=[], y=[], step=0.001, pot=11, **kwargs):
        """
        这是最小二乘拟合。
        :param dispersed:列表函数还是连续函数，默认列表函数，若为连续函数则需在kwargs中输入相关表达式。
        :param polynomial: 是否为多项式拟合，选False时激活非多项式拟合的基函数的使用。
        :param x: 传入拟合数据x，列表
        :param y: 传入拟合数据y，列表
        :param step: 拟合曲线的步长，一般情况无需调整，可以调大观察区别，一般拟合首位不适配就需要调整此参数
        :param kwargs: 用于传递各种参数，多项式拟合的最高次幂：highest_power=？
                                         非多项式拟合的基函数：fai=[?, ?, ...]
                                         fai内只能是基本初等函数，且线性无关，可以是(c为已知数字,x为变量,为自己所定义的变量符号)：
                                         输入的x必须为np.ndarray类型，且shape=（1，n）
                                         x**c,c**x,np.log(x)或np.log10(x)及其组合的其他底数的对数函数,
                                         np.sin(x),np.cos(x)……
                                         np.arcsin(x) ……
                                         需要常数项c时输入np.ones((1, len(x)))
                                         若想在结果中得到表达式、拟合图像、误差等需要输入自己的字符串型的基本初等函数，如：
                                         expression = ['1', 'x**2', 'np.log(x)'……]
                                         例子：
        LSF(polynomial=False, x=X, y=Y,
               fai=[np.ones((1, len(X))), np.array([X])**2, np.log(np.array([X])), np.sin(np.array([X])),
                    np.exp(np.array([X]))],
               expression=['1', 'x**2', 'np.log(x)', 'np.sin(x)', 'np.exp(x)'])
                                         persed_expression=“y=np.exp(x) + 1”，连续函数表达式,未知数统一用x表达。
                                         persed_ab = [a,b],拟合及绘制范围
                                         persed_step = 1000,拟合区间插值节点数。默认1000
        """
        self.dispersed = dispersed
        self.pot = pot
        self.polynomial = polynomial
        if self.dispersed == False:
            try:
                self.persed_expression = kwargs["persed_expression"].split("y=")[1]
                self.persed_ab = kwargs["persed_ab"]
                self.persed_step = kwargs["persed_step"]
                if self.persed_ab[0] < 0:
                    a = self.persed_ab[0] * 1.01
                else:
                    a = self.persed_ab[0] * 0.99
                if self.persed_ab[1] < 0:
                    b = self.persed_ab[1] * 0.99
                else:
                    b = self.persed_ab[1] * 1.01
                self.x = np.linspace(a, b, self.persed_step)
                self.x = self.x.reshape((self.x.shape[0], 1))
                self.y = np.array([[eval(self.persed_expression) for x in self.x]]).T
                self.y = self.y.reshape((self.y.shape[1], 1))
                print(self.x.shape)
                print(self.y.shape)
            except:
                print("连续函数表达式f(x)未定义！请检查输入参数！")
                sys.exit(1)
        else:
            self.x = np.array([x]).T
            self.y = np.array([y]).T
        try:
            self.n = kwargs["highest_power"]
        except:
            pass
        try:
            self.expression = kwargs["expression"]
        except:
            pass
        if polynomial == True:
            try:
                self.Polynomial(kwargs["highest_power"])
                self.painter(my_step=step)
            except:
                print("请检查输入是否有误？")
        else:
            try:
                self.Non_polynomial(kwargs["fai"])
                self.painter(my_step=step)
            except:
                print("请检查输入是否有误？")


    def Polynomial(self, n):
        """
        最小二乘拟合（多项式拟合）系数计算。
        :param n:
        :return:
        """
        m = self.x.shape[0]
        self.G = np.zeros((m, n+1))
        for i in range(m):
            for j in range(n+1):
                self.G[i][j] = (self.x[i][0])**j
        self.c = np.linalg.inv(self.G.T @ self.G) @ self.G.T @ self.y
        print("系数c：", list(self.c.reshape((n+1))))
        c = list(self.c.reshape((n+1)))
        result = ''
        for i in range(n):
            if c[i+1] >= 0:
                result = result + str(np.around(c[i], self.pot)) + "*x^" + str(i) + "+"
            else:
                result = result + str(np.around(c[i], self.pot)) + "*x^" + str(i)
        self.result = "y=" + result + str(np.around(c[-1], self.pot)) + "*x^" + str(len(c)-1)
        print("表达式：", self.result)

    def painter(self, my_step):
        """
        求误差，绘制曲线拟合图并保存结果！
        :param my_step:拟合时的步长。
        :return:
        """
        self.paint = plt.figure(1, figsize=(10, 5))
        plt.rcParams["font.sans-serif"] = ["SimHei"]
        plt.rcParams["axes.unicode_minus"] = False
        if self.dispersed==True:
            plt.scatter(self.x, self.y, marker="o", c="red")
        else:
            # pass
            plt.plot(self.x, self.y, c="green", linewidth=3, label="原函数曲线")
        plt.grid()
        plt.xlabel("x", fontsize=15)
        plt.ylabel("y", fontsize=15)
        if self.polynomial==True:
            plt.title(self.result)
        else:
            plt.title(self.result_show)
        if self.x[0][0]<0:
            a = self.x[0][0]*1.01
        else:
            a = self.x[0][0]*0.99
        if self.x[-1][0]<0:
            b = self.x[-1][0]*0.99
        else:
            b = self.x[-1][0]*1.01
        xx = np.arange(a, b, step=my_step)
        xx = xx.reshape((1, xx.shape[0]))
        try:
            #求解曲线拟合
            yy = []
            if self.polynomial==True:
                c = list(self.c.reshape((self.n + 1)))
                for j in range(xx.shape[1]):
                    result = 0
                    for i in range(self.n + 1):
                        result = result + c[i] * (xx[0][j]) ** i
                    yy.append(result)
                yy = np.array([yy])
            else:
                for j in range(xx.shape[1]):
                    result = self.caculate(xx[0][j])
                    yy.append(result)
                yy = np.array([yy])
            plt.plot(xx[0], yy[0], label='最小二乘拟合曲线')
        except:
            print("拟合出错！")

        #误差计算
        c_y = []
        for k in self.x:
            c_y.append(self.caculate(k))
        c_y = np.array(c_y)
        self.p_f = np.around((((c_y - self.y) ** 2).sum()) ** (0.5), self.pot)
        print("误差：", self.p_f)
        if self.dispersed==False:
            plt.legend()
        plt.show()

    def caculate(self, x):
        """
        通过拟合函数计算y值。
        :param x:计算的x值。
        :return:
        """
        if self.polynomial==True:
            c = list(self.c.reshape((self.n + 1)))
            result = 0
            for i in range(self.n + 1):
                result = result + np.around(c[i] * (x) ** i, self.pot)
        else:
            result = np.around(eval(self.result.replace("=", "").replace("y", "")), self.pot)
        return result

    def Non_polynomial(self, fai):
        self.fai = fai
        n = len(fai)
        m = self.x.shape[0]
        self.G = np.zeros((m, n))
        for i in range(m):
            for j in range(n):
                self.G[i][j] = fai[j][0][i]
        self.c = np.linalg.inv(self.G.T @ self.G) @ self.G.T @ self.y
        print("系数c：", list(self.c.reshape((n))))
        try:
            c = list(self.c.reshape((n)))
            result = ''
            for i in range(n-1):
                if c[i + 1] >= 0:
                    result = result + str(np.around(c[i], self.pot)) + "*" + self.expression[i] + "+"
                else:
                    result = result + str(np.around(c[i], self.pot)) + "*" + self.expression[i]
            self.result = "y=" + result + str(np.around(c[-1], self.pot)) + "*" + self.expression[-1]
            self.result_show = self.result.replace("**", "^")
            self.result_show = self.result_show.replace("np.log", "ln")
            self.result_show = self.result_show.replace("np.log10", "log")
            self.result_show = self.result_show.replace("np.exp", "e^")
            self.result_show = self.result_show.replace("np.sin", "sin")
            self.result_show = self.result_show.replace("np.cos", "cos")
            self.result_show = self.result_show.replace("np.tan", "tan")
            self.result_show = self.result_show.replace("np.arcsin", "arcsin")
            self.result_show = self.result_show.replace("np.arccos", "arccos")
            self.result_show = self.result_show.replace("np.arctan", "arctan")
            print("表达式：", self.result_show)
        except:
            pass

    def save_data(self, path):
        """
        保存计算数据的函数。
        :param path:保存路径及文件名（无文件类型后缀）。
        :return:
        """
        self.paint.savefig(path + "(拟合图像).png")
        if self.dispersed==False:
            if self.polynomial == True:
                data = pd.DataFrame(data=["y="+self.persed_expression, self.result, self.p_f], columns=["值"],
                                    index=["表达式", "拟合表达式", "估计误差"])
            else:
                data = pd.DataFrame(data=["y="+self.persed_expression, self.result_show, self.p_f],
                                    columns=["值"],
                                    index=["表达式", "拟合表达式", "估计误差"])
            data.to_csv(path + "(数据).csv")
        else:
            if self.polynomial == True:
                data = pd.DataFrame(data=[str(list(self.x)), str(list(self.y)), self.result, self.p_f], columns=["值"],
                                    index=["x", "y", "拟合表达式", "估计误差"])
            else:
                data = pd.DataFrame(data=[str(list(self.x)), str(list(self.y)), self.result_show, self.p_f],
                                    columns=["值"],
                                    index=["x", "y", "拟合表达式", "估计误差"])
            data.to_csv(path + "(数据).csv")

def red_csv_data(path):
    """
    这是读取csv文件用于拟合的快捷函数。
    :param path: 文件路径和名称
    :return: [x,y]
    """
    data = pd.read_csv(path)
    data = np.array(data)
    return [data[:, 0], data[:, 1]]

if __name__ == '__main__':
    # 5.1数据,多项式拟合
    X = [i*0.1 for i in range(1, 10, 1)]
    Y = [5.1234, 5.3057, 5.5687, 5.9375, 6.4370, 7.0978, 7.9493, 9.0253, 10.3627]

    # print("多项式拟合：")
    # lsf1 = LSF(polynomial=True, x=X, y=Y, highest_power=4, pot=5)
    # lsf1.save_data("./data/最优平方逼近5_1")
    #
    # print("**"*20)
    # # #5.1数据,非多项式拟合
    # print("非多项式拟合：")
    # lsf2 = LSF(polynomial=False, x=X, y=Y,
    #            fai=[np.ones((1, len(X))), np.array([X])**2, np.log(np.array([X])), np.sin(np.array([X])),
    #                 np.exp(np.array([X]))],
    #            expression=['1', 'x**2', 'np.log(x)', 'np.sin(x)', 'np.exp(x)'],
    #            pot=5)
    # lsf2.save_data("./data/最优平方逼近5_1_非多项式拟合")
    #
    # # 自定义数据计算
    # print("**"*20)
    # print("自定义数据：y=100sin(x)/x")
    # print("多项式拟合")
    # X_a = np.linspace(-10, 10, num=20)
    # Y_a = 100 * np.sin(X_a) / X_a
    # lsf3 = LSF(polynomial=True, x=X_a, y=Y_a, highest_power=10, pot=11)
    # print("**" * 20)
    #
    # print("非多项式拟合")
    # lsf4 = LSF(polynomial=False, x=X_a, y=Y_a,
    #            fai=[np.ones((1, len(X_a))), np.array([X_a])**(2), np.exp(np.array([X_a]))],
    #            expression=['1', 'x**(2)', "np.exp(x)"],
    #            pot=11)
    #
    # print("**" * 20)
    # xy = red_csv_data("./data/zmj.csv")
    # lsf5 = LSF(polynomial=True, x=xy[0], y=xy[1], highest_power=10, pot=2)
    # print("x={}时结果是{}".format(1.88, lsf5.caculate(1.88)))
    # lsf5.save_data("./data/正态分布数据zmj拟合")
    #
    # #连续函数的拟合
    # print("**" * 20)
    # lsf6 = LSF(dispersed=False, polynomial=True, highest_power=10, pot=6, persed_expression="y=np.sin(x)",
    #            persed_ab=[-2*np.pi, 2*np.pi], persed_step=100)
    # lsf6.save_data("./data/连续函数sin(x)的最小二乘拟合")

    #误差分析\
    # e = []
    # for i in range(0, 30):
    #     lsf7 = LSF(polynomial=True, x=X, y=Y, highest_power=i, pot=5)
    #     e.append(lsf7.p_f)
    # highest_power = [i for i in range(0, 30)]
    # plt.figure(2, figsize=(10, 5))
    # plt.rcParams["font.sans-serif"] = ["SimHei"]
    # plt.rcParams["axes.unicode_minus"] = False
    # plt.plot(highest_power, e)
    # plt.scatter(highest_power, e, marker="o", c="c")
    # plt.xlabel("多项式次数")
    # plt.ylabel("估计误差值")
    # plt.title("多项式拟合最高次数与误差关系曲线")
    # plt.show()

    #饱和水
    print("**" * 20)
    xy = red_csv_data("./data/饱和蒸汽表.csv")
    lsf5 = LSF(polynomial=True, x=xy[0], y=xy[1], highest_power=5, pot=6)
    print("x={}时结果是{}".format(1.44, lsf5.caculate(1.44)))